
# -*- coding: UTF-8 -*-

from pymongo import MongoClient
from WindPy import *

import time
import datetime
import math

def writeDailyCloseToTxt(startDateStr, endDateStr = time.strftime('%Y-%m-%d'), fillType='F', DBIP='localhost', DBPort=27017, DBName = 'stocks', collectionPreName = 'daily_close', txtFileName = 'close.txt'):



    client = MongoClient(DBIP, DBPort)


    db = client[DBName]

    collectionName = collectionPreName + '_' + startDateStr + '-' + endDateStr + '_' + fillType

    collection_dailyclose = db[collectionName]

    dateList = []
    totalStockCloseList = []

    dateListRes = collection_dailyclose.find({'StockID': '000001.SH'})
    dateListDict = dateListRes[0]

    for (k, v) in dateListDict.items():
        dateStr = ''
        try:
            dateStr = datetime.datetime.strptime(k, '%Y-%m-%d').strftime('%Y%m%d')
        except:
            dateStr = k
        dateList.append(dateStr)
    dateList = dateList[2:]
    dateList.insert(0, 'close')
    totalStockCloseList.append(dateList)

    for stockIdRes in collection_dailyclose.find({}, {'StockID':1}):
        stockIdStr = stockIdRes['StockID']
        for stockDailyCloseResDict in collection_dailyclose.find({'StockID':stockIdStr}):
            oneStockCloseList = []
            for i in stockDailyCloseResDict:
                oneStockCloseList.append(stockDailyCloseResDict[i])
            oneStockCloseList = oneStockCloseList[1:]
            totalStockCloseList.append(oneStockCloseList)
    print('Finished retrieving close data from MongoDB...')
    f = open(txtFileName, 'w')

    txtFileRowNum = len(oneStockCloseList)
    txtFileColNum = len(totalStockCloseList)

    f.write(str(txtFileRowNum-1) + '\t' + str(txtFileColNum-1) + '\t')

    for firstRowLen in range(txtFileColNum-1):
        f.write('\t')
    f.write('\n')
    print('Finished writing text file header...')


    for rowIter in range(txtFileRowNum):
        for colIter in range(txtFileColNum):
            try:
                cellData = totalStockCloseList[colIter][rowIter]
                if math.isnan(float(cellData)):
                    cellData = 0
                f.write(str(cellData) + '\t')
            except ValueError:
                cellData = str(totalStockCloseList[colIter][rowIter])
                f.write(cellData + '\t')

        f.write('\n')
        print('Finished writing row number ' + str(rowIter) + ' to the text file....')

    f.close()
    print('Finished writing the whole text file!')


writeDailyCloseToTxt(startDateStr='2012-01-01', collectionPreName='daily_close')